import torch
from model.Unet_model import Unet_model,Csp_unet_model
from model.utils import DownSampling, ConvBlock, UpSampling
from dataloader.dataload import Trans
from args import get_args
import cv2
import logging
import os
from model.csp_module import C3, Conv, dark_module
logging.basicConfig(level=logging.INFO)
import numpy as np


def predict(args, img_name, model):
    img = cv2.imread(img_name)
    img1 = img.copy()
    Transform = Trans()
    processed_img = Transform(img).to(args.device)
    output = model(processed_img.unsqueeze(0))
    output = torch.softmax(output, dim=1)
    output = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
    img1[output == 1] += np.array([80, 0, 0]).astype(np.uint8)
    img1[output == 2] += np.array([0, 0, 80]).astype(np.uint8)
    img1[output == 0] += np.array([0, 40, 0]).astype(np.uint8)
    cv2.imwrite(os.path.join(args.output, img_name.split('\\')[-1]), img1)


if __name__ == '__main__':
    args = get_args()
    logging.info("start to predict......")
    # model = Unet_model(ConvBlock, DownSampling, UpSampling).to(args.device)
    model = Csp_unet_model(ConvBlock, C3, Conv, UpSampling, base_channels=64, base_depth=1).to(args.device)
    model.load_state_dict(torch.load(args.ckpt_path))
    # model.eval()
    with torch.no_grad():
        for file in os.listdir(args.input):
            img_name = os.path.join(args.input, file)
            print(img_name)
            predict(args, img_name, model)
    logging.info("predict process is end,the result is stored in output directory")
